import { ChromaClient } from "chromadb";
import axios from "axios";
import { ElMessage } from "element-plus";
const chroma = new ChromaClient({ path:import.meta.env.VITE_AI_CHROMADB_API });
// 注意：chroma上线后httpVersion必须大于1.0，否则会报错，可以使用nginx反向代理，配置参考如下：
// location /chromadb {
//     rewrite ^/chromadb/(.*) /$1 break;
//     proxy_http_version 1.1;
//     proxy_pass http://localhost:8000;
// }

// 常用方法：
// chroma.getCollection({name: "my_collection"}) // 获取集合
// chroma.createCollection({name: "my_collection"}) // 创建集合
// chroma.deleteCollection({name: "my_collection"}) // 删除集合
// chroma.listCollections() // 获取所有集合
// collection.add(data) // 向集合中添加数据
// collection.delete(data) // 从集合中删除数据
// collection.count() // 查询集合中数据数量
// collection.peek() // 查询集合中的数据前10个
// collection.get() // 查询集合中的所有数据
// collection.query(queryData) // 查询集合中的数据
// collection.get({
// 	ids: ["id1", "id2", "id3", ...], //ids
// 	where: {"style": "style1"} // where
// }) //检索集合中的数据

class MyEmbeddingFunction {
    async generate(text) {
        let query = {
            "model": 'nomic-embed-text',
            "prompt": text[0],
        }
        let response = await axios.post(import.meta.env.VITE_AI_API_CHAT_EMB_URL, query)
        return response.data.embedding
    }
}

class ChromaDB {
    tableName = null;
    table = null;
    chroma = chroma;
    constructor(tableName) {
        this.tableName = tableName;
    }
    // 获得集合（有则获取无则创建）
    useTable = async (tableName) => {
        try {
            this.table = await chroma.getCollection({name: tableName, embeddingFunction:new MyEmbeddingFunction() })
        } catch (error) {
            this.table = await chroma.createCollection({ name: tableName, metadata: { "hnsw:space": "cosine" }, embeddingFunction:new MyEmbeddingFunction() });
            ElMessage({
                type: 'success',
                message: `添加成功`,
            })
        }
    }
    // 删除集合
    deleteTable = async (tableName) => {
        chroma.deleteCollection({name: tableName})
    }
    // 将文本转为向量
    textToEmbedding = async (text) => {
        let query = {
            "model": 'nomic-embed-text',
            "prompt": text,
        }
        let response = await axios.post(import.meta.env.VITE_AI_API_CHAT_EMB_URL, query)
        return response.data.embedding
    }
    // 向集合中添加数据(也可以直接在外部使用集合的add方法)
    add = async (ids,documents) => {
        let embeddings = []
        for(const i of documents){
            embeddings.push(await this.textToEmbedding(i))
        }
        let query = {ids,embeddings,documents}
        console.log(query)
        await this.table.add(query)
        ElMessage({
            type: "success",
            message: "上传成功"
        })
    }
    // 查询数据
    get = async (text,nResults=1) => {
        let embedding = await this.textToEmbedding(text)
        let result = await this.table.query({
            queryEmbeddings: embedding,
            nResults,
        });
        let docs = result.ids[0].map((item, index) => `#### 文章序号：${index+1} \n #### 文章标题：《${item}》\n #### 文章内容（文章内容位于{{{...}}}中）：{{{ ${result.documents[0][index]} }}} \n\n`);
        return docs;
    }
    // 从集合中删除数据
    deleteIds = async (ids) => {
        await this.table.delete({ids})
        ElMessage({
            type: "success",
            message: "删除成功"
        })
    }
    // 查询集合中数据数量
    count = async () => {
        return await this.table.count()
    }
    // 查询集合所有数据
    all = async () => {
        return await this.table.get()
    }
}


export {
    ChromaDB,
    chroma
};